//
// Created by hyh on 4/28/20.
//

#include <VideoRecord.h>

void VideoRecord::init(const char *destUrl, int w, int h) {

    avcodec_register_all();
    f = fopen(destUrl, "wb");
    if (!f) {
        LOGE("Could not open %s\n", destUrl);
        exit(1);
    }
    videoCodec = avcodec_find_encoder(AV_CODEC_ID_H264);
    if(!videoCodec){
        LOGE("vodeoCode %d not found",AV_CODEC_ID_H264);
        return;
    }

    videoContext = avcodec_alloc_context3(videoCodec);
    if(!videoContext){
        LOGE("Could not allocate video codec context");
        return;
    }

    videoPacket = av_packet_alloc();
    if(!videoPacket){
        LOGE("Could not allocate packet");
    }

    videoContext->bit_rate = 400000;
    videoContext->width = w;
    videoContext->height = h;
    videoContext->time_base = (AVRational){1,30};
    videoContext->framerate = (AVRational){30,1};
    videoContext->gop_size = 10;
    videoContext->max_b_frames = 1;
    videoContext->pix_fmt = AV_PIX_FMT_YUV420P;
    if(videoCodec->id==AV_CODEC_ID_H264){
        av_opt_set(videoContext->priv_data, "preset", "superfast", 0);
    }
    int ret = avcodec_open2(videoContext,videoCodec,NULL);
    if(ret<0){
        LOGE("Could not open codec:%s",av_err2str(ret));
        return;
    }

    videoFrame = av_frame_alloc();
    if(!videoFrame){
        LOGE("Could not allocate video frame");
        return;
    }

    videoFrame->format = videoContext->pix_fmt;
    videoFrame->width = videoContext->width;
    videoFrame->height = videoContext->height;

    ret = av_frame_get_buffer(videoFrame,32);
    if(ret<0){
        LOGE("Could not allocate the video frame data");
        return;
    }

    ret = av_frame_make_writable(videoFrame);
    if (ret < 0)
        exit(1);
}


void VideoRecord::encode(uint8_t *data,int len) {


    int size = videoContext->width*videoContext->height;
    memcpy(videoFrame->data[0],data,size);
    for (int j = 0; j < size / 4; j++) {
        *(videoFrame->data[2] + j) = *(data + size + j * 2);
        // V
        *(videoFrame->data[1] + j) = *(data + size + j * 2 + 1);
        //U
    }
    videoFrame->pts=video_count++;
    /* encode the image */
    encode(videoContext, videoFrame, videoPacket, f);
}


void VideoRecord::encode(AVCodecContext *enc_ctx, AVFrame *frame, AVPacket *pkt, FILE *outfile) {
    int ret;
    /* send the frame to the encoder */
    if (frame)
        LOGE("Send frame %lld", frame->pts);
    ret = avcodec_send_frame(enc_ctx, frame);
    if (ret < 0) {
        LOGE("Error sending a frame for encoding\n");
        exit(1);
    }
    while (ret >= 0) {
        ret = avcodec_receive_packet(enc_ctx, pkt);
        if (ret == AVERROR(EAGAIN) || ret == AVERROR_EOF)
            return;
        else if (ret < 0) {
            LOGE("Error during encoding\n");
            exit(1);
        }
        LOGE("Write packet %lld (size=%5d)\n", pkt->pts, pkt->size);
        fwrite(pkt->data, 1, pkt->size, outfile);
        av_packet_unref(pkt);
    }
}

void VideoRecord::stop() {

    uint8_t endcode[] = { 0, 0, 1, 0xb7 };
    /* flush the encoder */
    encode(videoContext, NULL, videoPacket, f);
    /* add sequence end code to have a real MPEG file */
//    fwrite(endcode, 1, sizeof(endcode), f);
    fclose(f);
    avcodec_free_context(&videoContext);
    av_frame_free(&videoFrame);
    av_packet_free(&videoPacket);
}